Fix Utf8JsonReader invalid-literal errors: unbounded message, dropped bytes, IndexOutOfRangeException - #134448

using BenchmarkDotNet.Running;
using System;
using System.Buffers;
using System.Text;
using System.Text.Json;
using BenchmarkDotNet.Attributes;

BenchmarkSwitcher.FromAssembly(typeof(LiteralMultiSegmentBenchmarks).Assembly).Run(args);

[MemoryDiagnoser]
public class LiteralMultiSegmentBenchmarks
{
    private const int ItemCount = 1000;
    private byte[] _utf8 = null!;
    private ReadOnlySequence<byte> _sequence;

    public enum PayloadKind { Literals, Objects }
    public enum InputLayout { Span, SingleSegmentSequence, Bytes1, Bytes2, Bytes3, Bytes64, Bytes4096 }

    [ParamsAllValues]
    public PayloadKind Payload { get; set; }

    [ParamsAllValues]
    public InputLayout Layout { get; set; }

    [GlobalSetup]
    public void Setup()
    {
        var builder = new StringBuilder("[");
        for (int i = 0; i < ItemCount; i++)
        {
            if (i != 0)
            {
                builder.Append(", ");
            }

            if (Payload == PayloadKind.Objects)
            {
                builder.Append(' ', i % 4);
            }

            // The 17-byte literal-group stride varies alignment at 64-byte and 4 KiB boundaries.
            builder.Append(Payload == PayloadKind.Literals
                ? "true,false,null"
                : FormattableString.Invariant($"{{\"id\":{i},\"active\":{(i % 2 == 0 ? "true" : "false")},\"parent\":null,\"name\":\"item {i}\"}}"));
        }

        builder.Append(']');
        _utf8 = Encoding.UTF8.GetBytes(builder.ToString());
        int segmentSize = Layout switch
        {
            InputLayout.Span or InputLayout.SingleSegmentSequence => _utf8.Length,
            InputLayout.Bytes1 => 1,
            InputLayout.Bytes2 => 2,
            InputLayout.Bytes3 => 3,
            InputLayout.Bytes64 => 64,
            InputLayout.Bytes4096 => 4096,
            _ => throw new InvalidOperationException($"Unexpected layout: {Layout}.")
        };

        if (segmentSize == _utf8.Length)
        {
            _sequence = new ReadOnlySequence<byte>(_utf8);
        }
        else
        {
            var first = new Segment(_utf8.AsMemory(0, segmentSize));
            Segment last = first;
            for (int offset = segmentSize; offset < _utf8.Length; offset += segmentSize)
            {
                last = last.Append(_utf8.AsMemory(offset, Math.Min(segmentSize, _utf8.Length - offset)));
            }

            _sequence = new ReadOnlySequence<byte>(first, 0, last, last.Memory.Length);
        }

        var reader = Layout == InputLayout.Span ? new Utf8JsonReader(_utf8) : new Utf8JsonReader(_sequence);
        int tokens = 0;
        int literals = 0;
        int splitLiterals = 0;
        while (reader.Read())
        {
            tokens++;
            if (reader.TokenType is JsonTokenType.True or JsonTokenType.False or JsonTokenType.Null)
            {
                literals++;
                if (reader.HasValueSequence)
                {
                    splitLiterals++;
                }
            }
        }

        int expectedTokens = Payload == PayloadKind.Literals ? 3 * ItemCount + 2 : 10 * ItemCount + 2;
        int expectedLiterals = Payload == PayloadKind.Literals ? 3 * ItemCount : 2 * ItemCount;
        bool multiSegment = segmentSize < _utf8.Length;
        if (reader.BytesConsumed != _utf8.Length || tokens != expectedTokens || literals != expectedLiterals
            || _sequence.IsSingleSegment == multiSegment || (multiSegment && splitLiterals == 0))
        {
            throw new InvalidOperationException(
                $"Invalid coverage for {Payload}/{Layout}: {reader.BytesConsumed} bytes, {tokens} tokens, {literals} literals, {splitLiterals} split literals.");
        }
    }

    [Benchmark]
    public long Read()
    {
        var reader = Layout == InputLayout.Span ? new Utf8JsonReader(_utf8) : new Utf8JsonReader(_sequence);
        while (reader.Read())
        {
        }

        return reader.BytesConsumed;
    }

    private sealed class Segment : ReadOnlySequenceSegment<byte>
    {
        public Segment(ReadOnlyMemory<byte> memory) => Memory = memory;

        public Segment Append(ReadOnlyMemory<byte> memory)
        {
            var next = new Segment(memory) { RunningIndex = RunningIndex + Memory.Length };
            Next = next;
            return next;
        }
    }
}